import encoding
import threading
import SocketServer
import socket
import struct
import os
from Queue import Queue, Empty
from Crypto.Random import random
from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA
from Crypto.Hash import RIPEMD
from Crypto.Hash import HMAC
from enigma import EnigmaMachine, add_uppercase

listen_at = ('127.0.0.1', 4500)
server_addr = ('127.0.0.1', 80)
edge_conn_timeout = 20
RSA_keyfile = '.\pubkey.key'
enigma_settings = {'count': 2, 'rotors': 'II VI V', 'ring_settings': "1 16 7", 'reflector': 'B'}
complete_alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
HTTP_s_header = "Host: " + server_addr[0] + "\r\n"\
"Connection: keep-alive\r\n"\
"User-Agent: Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/40.0.2214.115 Safari/537.36\r\n"\
"Accept-Encoding: gzip, deflate, lzma\r\n"\
"\r\n"
HTTP_header = "Host: " + server_addr[0] + "\r\n"\
"Connection: keep-alive\r\n"\
"Content-Length: %s\r\n"\
"Content-Type: application/x-www-form-urlencoded\r\n"\
"Accept: */*\r\n"\
"User-Agent: Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/40.0.2214.115 Safari/537.36\r\n"\
"Accept-Encoding: gzip, deflate, lzma\r\n"\
"\r\n"
server_HTTP_header = "HTTP/1.1 200 OK\r\n"\
"Server: Apache\r\n"\
"Vary: Accept-Encoding\r\n"\
"Transfer-Encoding: chunked\r\n"\
"Content-Type: text/html; charset=utf-8\r\n"\
"Connection: keep-alive\r\n"\
"Cache-Control: private\r\n"
server_header_list = frozenset(server_HTTP_header.split('\r\n'))

class Logger(object):
    def __init__(self):
        self.lock = threading.Lock()

    def log(self, str):
        self.lock.acquire()
        print(str)
        self.lock.release()

logger = Logger()

logger.log('Starting...')

EnigmaMachine.wrapper_initialize(**enigma_settings)

try:
    server_socket = socket.create_connection(server_addr)
    server_socket.send('GET /' + encoding.url_add_slashes_underline(encoding.create_random_string(17, complete_alphabet)) + ' HTTP/1.1\r\n' + HTTP_s_header)
    sresp = server_socket.recv(4096).split('\r\n\r\n')
    if not sresp[0][0:15] == 'HTTP/1.1 200 OK':
        raise
    seedstr = sresp[1][0:4]
    if not seedstr:
        raise
except:
    logger.log('Failed to connect to server.')
    os._exit(0)

enigmakey =''.join(random.sample(complete_alphabet, 3))
hmackey = struct.pack("!I", random.getrandbits(32))

alphabet = ''.join(random.sample(complete_alphabet, 24))

with open(RSA_keyfile, 'r') as RSA_key:
    key = RSA.importKey(RSA_key.read())
    rsa = PKCS1_OAEP.new(key)
padding_count = random.randint(0, 16)
padding = ''
for i in xrange(padding_count):
    padding += struct.pack("!I", random.getrandbits(32))
hs = rsa.encrypt(seedstr + enigmakey + alphabet + hmackey) + padding
server_socket.send('POST /' + encoding.url_add_slashes_underline(encoding.create_random_string(17, complete_alphabet)) + ' HTTP/1.1\r\n' + (HTTP_header % len(hs)) + hs)

upstream_hmac = HMAC.new(hmackey, digestmod=RIPEMD.new()) # used to sign data waiting to be sent
downstream_hmac = HMAC.new(hmackey, digestmod=RIPEMD.new()) # used to verify data we have received
upstream_enigma = EnigmaMachine() # used to encrypt data waiting to be sent
upstream_enigma.set_display(enigmakey)
downstream_enigma = EnigmaMachine() # used to decrypt data we have received
downstream_enigma.set_display(enigmakey)

try:
    vf = server_socket.recv(4096).split('\r\n\r\n')[1].replace(' ', '').upper()
    _alphabet = upstream_enigma.process_text(vf)[0:24]
except:
    logger.log('Received bad handshake data from server.')
    os._exit(0)

if not _alphabet == alphabet:
    logger.log('Enigma machine desync.')
    os._exit(0)

logger.log('Connected!')

outbuf = Queue()
clean_conn = Queue()
channels = []

class SenderThread(threading.Thread):
    def run(self):
        while True:
            msg = outbuf.get()
            downstream_hmac.update(msg)
            hmac = downstream_hmac.digest()

            length = struct.pack("!I", len(hmac) + len(msg))

            encdata = add_uppercase(downstream_enigma.process_text(encoding.enbase(random.choice(alphabet) + length + hmac + msg, alphabet)))

            url_len = random.randint(0, min(48, len(encdata)))
            url, body = encoding.url_add_slashes_underline(encdata[0:url_len]), encoding.add_space(encdata[url_len:])

            data = 'POST /' + url + ' HTTP/1.1\r\n' + (HTTP_header % len(body)) + body

            try:
                server_socket.sendall(data)
            except:
                logger.log('Connection closed.')
                os._exit(0)

class ReceiverThread(threading.Thread):
    def run(self):
        inbuf = []
        ib = ''
        while True:
            data = server_socket.recv(4096)
            if not data:
                logger.log('Connection closed.')
                os._exit(0)

            while True:
                data = ib + data
                if not data:
                    break
                if '\r\n' in data:
                    d = data.split('\r\n', 1)
                    if d[0] in server_header_list or not d[0] or d[0] == '0':
                        data = d[1]
                        continue
                    length = int(d[0], 16)
                    if len(d[1]) < length:
                        ib = data
                        break
                    else:
                        ib = ''
                        inbuf.append(d[1][:length])
                        data = d[1][length+2:]
                else:
                    ib = data

            if inbuf == []:
                continue

            for data in inbuf:
                try:
                    msg = encoding.debase(upstream_enigma.process_text(data.replace(' ', '')), alphabet)
                except:
                    logger.log('Exception decoding data from server.')
                    os._exit(0)

                if not msg:
                    break

                hmac = msg[1:21]
                upstream_hmac.update(msg[21:])

                if not upstream_hmac.digest() == hmac:
                    logger.log('Got a message with bad HMAC.')
                    os._exit(0)

                if msg[21] == 'c':
                    clean_conn.put(msg[22:26])
                elif msg[21] == 'd':
                    for i in channels:
                        if i.connected(msg[22:26]):
                            i.disconnect()
                            break
                elif msg[21] == 'r':
                    for i in channels:
                        if i.connected(msg[22:26]):
                            i.push(msg[26:])
                            break
                elif msg[21] == 'p':
                    pass
                else:
                    logger.log('Connection closed by server.')
                    os._exit(0)

            inbuf = []

class RequestHandler(SocketServer.StreamRequestHandler):
    def handle(self):
        self.disconnected = False
        outbuf.put('c')

        try:
            self.id = clean_conn.get(True, edge_conn_timeout)
        except Empty:
            logger.log('Failed to get a connection from server.')
            return

        self.in_queue = Queue()
        channels.append(self)
        self.sender_thread = threading.Thread(target=self.send_to_req)
        self.sender_thread.start()

        while True:
            if self.disconnected:
                return
            data = self.request.recv(65536)
            if not data:
                outbuf.put('d' + self.id)
                return
            outbuf.put('m' + self.id + data)

    def connected(self, num):
        if self.disconnected or not self.id:
            return False
        else:
            return self.id == num

    def disconnect(self):
        self.disconnected = True

    def push(self, msg):
        self.in_queue.put(msg)

    def send_to_req(self):
        while True:
            msg = self.in_queue.get()
            try:
                self.request.sendall(msg)
            except:
                outbuf.put('d' + self.id)
                self.disconnected = True
                break

if __name__ == '__main__':
    SenderThread().start()
    ReceiverThread().start()
    SocketServer.ThreadingTCPServer(listen_at, RequestHandler).serve_forever()